package org.ala.tiktools.sharddb.db.provider;

import javax.sql.DataSource;
import java.io.PrintWriter;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.SQLFeatureNotSupportedException;
import java.util.concurrent.Callable;
import java.util.logging.Logger;

/**
 * 流量分流代理数据源
 * <p>  适用于流量实在太高的场景
 * <p>  单个数据源连接池是通过多个线程共同竞争资源，来实现链接的借出
 * <p>  当流量较大时，借出接口的阻塞会非常严重
 * <p>  这里给原数据源包了一层，准备多份同样的数据源
 * <p>      当流量较大时，用 thread.id + 时间戳 的方式 hash 到其中一个真实数据源，从而实现分流目的
 *
 * @author ala
 * @date 2025-03-25 12:36
 */
public class DiversionDataSourceProxy implements DataSource {


    /**
     *  代理数据源大小
     */
    protected int size;
    /**
     *  代理数据源数组
     */
    protected DataSource[] dataSources;
    /**
     *  构建数据源
     */
    protected Callable<DataSource> initer;

    public static DiversionDataSourceProxy build(int size, boolean initPool, Callable<DataSource> initer) {
        DiversionDataSourceProxy p = new DiversionDataSourceProxy();
        p.size = size;
        p.initer = initer;
        p.init(initPool);
        return p;
    }
    public void init(boolean initPool) {
        dataSources = new DataSource[size];

        //  是否需要初始化每一个池子
        if (initPool) {
            for (int i = 0; i < size; i++) {
                getDataSource(i);
            }
        }
    }

    /**
     * 初始化第i个连接池
     * <p>  i不能超过256
     */
    protected DataSource getDataSource(int i) {
        if (dataSources[i] == null) {
            final Object lock = i < 128 ? i : Integer.toString(i).intern();
            synchronized (lock) {
                if (dataSources[i] == null) {
                    try {
                        dataSources[i] = this.initer.call();
                    } catch (Exception e) {
                        throw new RuntimeException("初始化连接池异常", e);
                    }
                }
            }
        }
        return dataSources[i];
    }
    protected DataSource getDataSource() {
        long id = Thread.currentThread().getId(), now = System.currentTimeMillis();
        return getDataSource((int)((id + now) % dataSources.length));
    }


    @Override
    public Connection getConnection() throws SQLException {
        return getDataSource().getConnection();
    }
    @Override
    public Connection getConnection(String username, String password) throws SQLException {
        return getDataSource().getConnection(username, password);
    }

    @Override
    public <T> T unwrap(Class<T> iface) throws SQLException {
        return getDataSource().unwrap(iface);
    }

    @Override
    public boolean isWrapperFor(Class<?> iface) throws SQLException {
        return getDataSource().isWrapperFor(iface);
    }

    @Override
    public PrintWriter getLogWriter() throws SQLException {
        return getDataSource().getLogWriter();
    }

    @Override
    public void setLogWriter(PrintWriter out) throws SQLException {
        for (int i = 0; i < dataSources.length; i++) {
            getDataSource(i).setLogWriter(out);
        }
    }

    @Override
    public void setLoginTimeout(int seconds) throws SQLException {
        for (int i = 0; i < dataSources.length; i++) {
            getDataSource(i).setLoginTimeout(seconds);
        }
    }

    @Override
    public int getLoginTimeout() throws SQLException {
        return getDataSource().getLoginTimeout();
    }

    @Override
    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
        return getDataSource().getParentLogger();
    }
}
